{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "\n",
    "# 定义 PINN 神经网络\n",
    "class PINN(nn.Module):\n",
    "    def __init__(self, input_dim, hidden_dim, output_dim):\n",
    "        super(PINN, self).__init__()\n",
    "        self.model = nn.Sequential(\n",
    "            nn.Linear(input_dim, hidden_dim),\n",
    "            nn.Tanh(),\n",
    "            nn.Linear(hidden_dim, hidden_dim),\n",
    "            nn.Tanh(),\n",
    "            nn.Linear(hidden_dim, output_dim)\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.model(x)\n",
    "\n",
    "# 定义 HJB 方程的损失函数\n",
    "def hjb_loss(model, x):\n",
    "    x.requires_grad = True\n",
    "    u = model(x)\n",
    "\n",
    "    # 计算梯度\n",
    "    u_x = torch.autograd.grad(u.sum(), x, create_graph=True)[0]\n",
    "    u_xx = torch.autograd.grad(u_x.sum(), x, create_graph=True)[0]\n",
    "\n",
    "    # HJB 方程残差\n",
    "    residual = u + u_x * x - 0.5 * u_xx\n",
    "    return torch.mean(residual ** 2)\n",
    "\n",
    "# 训练 PINN\n",
    "def train_pinn():\n",
    "    input_dim, hidden_dim, output_dim = 1, 50, 1\n",
    "    model = PINN(input_dim, hidden_dim, output_dim)\n",
    "    optimizer = optim.Adam(model.parameters(), lr=0.01)\n",
    "\n",
    "    # 生成训练数据\n",
    "    x_train = torch.linspace(-1, 1, 100).view(-1, 1)\n",
    "\n",
    "    for epoch in range(10000):\n",
    "        optimizer.zero_grad()\n",
    "        loss = hjb_loss(model, x_train)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "        if epoch % 1000 == 0:\n",
    "            print(f\"Epoch {epoch}, Loss: {loss.item()}\")\n",
    "\n",
    "    return model\n",
    "\n",
    "# 运行训练\n",
    "trained_model = train_pinn()\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.9.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
